#!/bin/bash

set -e

source "$(
  cd "$(dirname "$0")" || exit 1
  pwd
)/./common.sh"

#下载依赖和工具包
function install_development() {
  sudo yum install -y wget tar unzip patch cmake
  source /etc/profile
}

#获取TensorRT-OSS-8.2.0
function download_TensorRT-OSS() {
  echo "# TensorRT-OSS-8.2.0正在下载"
  #git clone https://github.com/NVIDIA/TensorRT.git "$1"/TensorRT-8.2.0
  cp -r  "$2"/src/TensorRT-8.2.0 "$1"/
  if [[ $? -eq 0 ]]; then
    echo -e "\033[1;32;1m# TensorRT-OSS-8.2.0克隆成功\033[0m"
  else
    echo -e "\033[31m# TensorRT-OSS-8.2.0克隆失败,请重新执行脚本\033[0m"
    exit 1
  fi

  echo "# 正在获取glibc安装包"
  cp "$2"/src/glibc-2.34-79.oe2203.aarch64.rpm "$1"/TensorRT-8.2.0
  cp "$2"/src/glibc-common-2.34-79.oe2203.aarch64.rpm "$1"/TensorRT-8.2.0
  cp "$2"/src/glibc-devel-2.34-79.oe2203.aarch64.rpm "$1"/TensorRT-8.2.0
  if [[ $? -eq 0 ]]; then
    echo -e "\033[1;32;1m# 获取glibc安装包成功\033[0m"
  else
    echo -e "\033[31m# 获取glibc安装包克隆失败,请重新执行脚本\033[0m"
    exit 1
  fi
}

function install_TensorRT-OSS() {
  echo "# TensorRT-OSS-8.2.0正在安装"
  cp "$3"/TensorRT.patch "$1"/TensorRT-8.2.0
  cd "$1"/TensorRT-8.2.0
  if [[ $? -eq 0 ]]; then
    echo -e "\033[1;32;1m# TensorRT-OSS-8.2.0子模块克隆成功\033[0m"
  else
    echo -e "\033[31m# TensorRT-OSS-8.2.0子模块克隆失败,请重新执行脚本\033[0m"
    exit 1
  fi
  patch -p1 <./TensorRT.patch
  sed -i "2 i\set (CMAKE_C_COMPILER  \"$2/bisheng-compiler-2.1.0/bin/clang\")" CMakeLists.txt
  sed -i "3 i\set (CMAKE_CXX_COMPILER  \"$2/bisheng-compiler-2.1.0/bin/clang++\")" CMakeLists.txt
  cp "$2"/TensorRT-8.2.0.6/include/NvInferConsistency.h  "$1"/TensorRT-8.2.0/include/
  cp "$2"/TensorRT-8.2.0.6/include/NvOnnxParser.h  "$1"/TensorRT-8.2.0/include/
  cp "$2"/TensorRT-8.2.0.6/include/NvInferConsistencyImpl.h  "$1"/TensorRT-8.2.0/include/
  cp "$2"/TensorRT-8.2.0.6/include/NvInferSafeRuntime.h  "$1"/TensorRT-8.2.0/include/

  echo "# glibc-2.34正在安装"
  rpm -Uvh glibc-2.34-79.oe2203.aarch64.rpm glibc-common-2.34-79.oe2203.aarch64.rpm glibc-devel-2.34-79.oe2203.aarch64.rpm
  if [[ $? -eq 0 ]]; then
    echo -e "\033[1;32;1m# glibc-2.34安装成功\033[0m"
  else
    echo -e "\033[31m# glibc-2.34安装失败,请重新执行脚本\033[0m"
    exit 1
  fi

  mkdir build && cd build
  export TRT_LIBPATH="$2"/TensorRT-8.2.0.6/targets/aarch64-linux-gnu
	export TRT_OSSPATH="$1"/TensorRT-8.2.0
  export TRT_OUT_DIR="$1"/TensorRT-8.2.0/build/out
  cmake .. -DGPU_ARCHS="80" -DTRT_LIB_DIR=$TRT_LIBPATH/lib/  -DTRT_OUT_DIR="$1"/TensorRT-8.2.0/build/out -DTRT_PLATFORM_ID=aarch64 -DCUDA_VERSION=11.4
  ln -s $TRT_OUT_DIR/libnvinfer_plugin.so $TRT_OUT_DIR/libnvinfer_plugin.so.8
  CC=clang CXX=clang++  make -j4
  make install 
  echo -e "\033[1;32;1m# TensorRT-OSS-8.2.0已经安装完毕\033[0m"
}

function write_TensorRT-GA_module() {
  cat >"$1/TensorRT-8.2.0.6/TensorRT_modulefiles" <<EOF
#%Module1.0
conflict TensorRT
variable modfile [file normalize [info script]]
proc getModulefileDir {} {
  variable modfile
  set modfile_path [file dirname \$modfile]
  return \$modfile_path
}
set pwd [getModulefileDir]
set TensorRT \$pwd
setenv TensorRT \$TensorRT
prepend-path PATH \$TensorRT/bin
prepend-path LD_LIBRARY_PATH \$TensorRT/targets/aarch64-linux-gnu/lib
EOF
}

function main() {
    install_development
    download_TensorRT-OSS "$1" "$3"
    install_TensorRT-OSS "$1" "$2" "$3" 
    write_TensorRT-GA_module "$2"
    #rm -rf "$1"/TensorRT-8.2.0
}

main "$@"